# -*- coding: utf-8 -*-
"""
轨迹数据采样器模块

该模块提供了轨迹数据的批量采样功能，用于强化学习训练过程中的小批量数据生成。
支持内存安全检查、数据重组和迭代器生成等功能。

主要功能：
1. 将多个轨迹合并成大的数据池
2. 支持批量采样和小批量训练
3. 处理嵌套数据结构（使用'>'分隔的键名）
4. 内存溢出保护和动态批次大小调整
"""

import math
import numpy as np
from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler


class TrajectoryDataSampler():
    """
    轨迹数据采样器类

    负责将多个轨迹数据组织成可采样的数据集，并提供批量采样功能。
    支持强化学习训练中常见的数据格式，包括状态、动作、奖励等。

    Attributes:
        num_batches (int): 批次数量，将大数据池分割成多少个小批次
        memory_safety_check (bool): 是否启用内存安全检查，防止GPU OOM
        container (dict): 数据容器，存储所有合并后的轨迹数据
        big_batch_size (int): 大批次大小（所有轨迹的总长度）
        mini_batch_size (int): 小批次大小（每个训练批次的大小）
    """

    def __init__(self, n_div, traj_pool, flag, memory_safety_check=False):
        """
        初始化轨迹数据采样器

        Args:
            n_div (int): 批次分割数量，将数据分成n_div个批次
            traj_pool (list): 轨迹池，包含多个轨迹对象的列表
            flag (str): 标识符，目前只支持'train'
            memory_safety_check (bool): 是否启用内存安全检查，默认False

        功能：
        1. 验证输入参数的合法性
        2. 检查轨迹数据是否包含必需的字段
        3. 将多个轨迹数据合并到容器中
        4. 处理嵌套数据结构
        5. 计算批次大小
        """
        self.num_batches = n_div  # 批次数量
        self.memory_safety_check = memory_safety_check  # 内存安全检查标志

        # 内存安全检查模式下，批次数量必须为1
        if self.memory_safety_check:
            assert self.num_batches == 1, "内存安全检查模式下批次数量必须为1"

        self.num_batch = None  # 批次编号（未使用）
        self.container = {}    # 数据容器，用于存储合并后的轨迹数据
        self.warned = False    # 警告标志（未使用）

        # 目前只支持训练模式
        assert flag == 'train', "目前只支持'train'模式"

        # 定义必需的数据字段（2024/08修改后的版本）
        # 这些是强化学习训练中基本的数据元素
        req_dict = ['avail_act', 'state', 'action', 'reward', 'next_state', 'done']
        req_dict_rename = ['avail_act', 'state', 'action', 'reward', 'next_state', 'done']

        # 处理嵌套数据结构（如 'obs>image', 'obs>vector' 等）
        # 如果直接字段不存在，查找嵌套字段并添加到需求列表中
        for key_index, key in enumerate(req_dict):
            key_name = req_dict[key_index]
            key_rename = req_dict_rename[key_index]

            # 检查轨迹中是否存在直接字段
            if not hasattr(traj_pool[0], key_name):
                # 查找嵌套字段（包含'>'的键名）
                real_key_list = [real_key for real_key in traj_pool[0].__dict__
                               if (key_name + '>' in real_key)]
                assert len(real_key_list) > 0, f"缺少必需的变量: {key}"

                # 将所有嵌套字段添加到需求列表
                for real_key in real_key_list:
                    mainkey, subkey = real_key.split('>')
                    req_dict.append(real_key)
                    req_dict_rename.append(key_rename + '>' + subkey)

        # 初始化大批次大小，用于验证数据一致性
        self.big_batch_size = -1

        # 将轨迹数据加载到容器中
        for key_index, key in enumerate(req_dict):
            key_name = req_dict[key_index]
            key_rename = req_dict_rename[key_index]

            # 跳过不存在的字段
            if not hasattr(traj_pool[0], key_name):
                continue

            # 将所有轨迹的相同字段数据合并成一个大的数组
            # 这是关键的数据合并步骤，将分散在多个轨迹中的数据按时间顺序连接
            set_item = np.concatenate([getattr(traj, key_name) for traj in traj_pool], axis=0)

            # 验证所有字段的数据长度一致性
            # 确保所有数据字段具有相同的时间步数量
            assert self.big_batch_size == set_item.shape[0] or (self.big_batch_size < 0), \
                f"数据长度不一致: {key}, 期望: {self.big_batch_size}, 实际: {set_item.shape[0]}"
            self.big_batch_size = set_item.shape[0]

            # 将合并后的数据存储到容器中
            self.container[key_rename] = set_item

        # 计算每个小批次的大小
        # 向上取整确保所有数据都能被采样到
        self.mini_batch_size = math.ceil(self.big_batch_size / self.num_batches)

        # 注释掉的代码：优势函数的标准化处理（原PPO算法中的步骤）
        # self.container[advantage_rename] = self.container[return_rename] - self.container[value_rename]
        # self.container[advantage_rename] = (self.container[advantage_rename] -
        #                                    self.container[advantage_rename].mean()) / \
        #                                   (self.container[advantage_rename].std() + 1e-5)

    def __len__(self):
        """
        返回批次数量的长度

        Returns:
            int: 批次数量
        """
        return self.num_batches

    def determine_max_n_sample(self):
        """
        确定最大安全采样数量

        在内存安全检查模式下，动态调整采样数量以避免GPU内存溢出。
        基于历史经验和当前批次大小来计算安全的采样数量。

        Returns:
            int: 最大安全采样数量

        工作原理：
        1. 首次使用时，初始化采样数量序列
        2. 如果没有发生过OOM，可以使用全部数据
        3. 如果发生过OOM，使用历史最大安全数量
        """
        assert self.memory_safety_check, "此方法仅在内存安全检查模式下可用"

        # 首次初始化：创建一个采样数量序列，从1%到100%递增
        if not hasattr(TrajectoryDataSampler, 'MaxSampleNum'):
            # 初始化：创建50个采样点，从1%到100%
            TrajectoryDataSampler.MaxSampleNum = [
                int(self.big_batch_size * (i + 1) / 50) for i in range(50)
            ]
            max_n_sample = self.big_batch_size

        # 如果最近记录的最大采样数量大于0，说明没有发生过OOM
        elif TrajectoryDataSampler.MaxSampleNum[-1] > 0:
            # 只有当批次大小增加时才更新记录
            if self.big_batch_size > TrajectoryDataSampler.MaxSampleNum[-1]:
                TrajectoryDataSampler.MaxSampleNum.append(self.big_batch_size)
            max_n_sample = self.big_batch_size

        # 如果最近记录为0，说明发生过OOM，使用倒数第二个记录值
        else:
            assert TrajectoryDataSampler.MaxSampleNum[-2] > 0, "采样数量记录异常"
            max_n_sample = TrajectoryDataSampler.MaxSampleNum[-2]

        return max_n_sample

    def reset_and_get_iter(self):
        """
        重置采样器并返回迭代器

        创建一个新的数据采样迭代器，支持内存安全检查和嵌套数据结构处理。
        每次调用都会重新随机化数据顺序。

        Yields:
            dict: 包含一个批次数据的字典，数据已经重新组织成嵌套结构

        功能：
        1. 创建批次采样器
        2. 内存安全检查下的采样数量调整
        3. 数据索引采样和批量提取
        4. 嵌套数据结构的重组
        """
        # 普通模式：使用全部数据进行随机采样
        if not self.memory_safety_check:
            # 创建随机子集采样器，然后按小批次大小进行批量采样
            self.sampler = BatchSampler(
                SubsetRandomSampler(range(self.big_batch_size)),
                self.mini_batch_size,
                drop_last=False
            )

        # 内存安全检查模式：限制采样数量
        else:
            max_n_sample = self.determine_max_n_sample()
            n_sample = min(self.big_batch_size, max_n_sample)

            # 首次提醒用户采样数量被限制的情况
            if not hasattr(self, 'reminded'):
                self.reminded = True
                drop_percent = (self.big_batch_size - n_sample) / self.big_batch_size * 100

                # 如果丢弃的数据比例过大，发出警告或报错
                if drop_percent > 20:
                    print(f'警告：丢弃了 {drop_percent:.1f}% 的样本数据！')
                    assert False, "GPU内存不足！请减少批次大小或启用内存优化。"
                else:
                    print(f'内存优化：丢弃了 {drop_percent:.1f}% 的样本数据')

            # 使用限制后的采样数量创建采样器
            self.sampler = BatchSampler(
                SubsetRandomSampler(range(n_sample)),
                n_sample,
                drop_last=False
            )

        # 生成批次数据
        for indices in self.sampler:
            selected = {}

            # 根据采样索引提取所有字段的数据
            for key in self.container:
                selected[key] = self.container[key][indices]

            # 处理嵌套数据结构，将 'mainkey>subkey' 重组为嵌套字典
            for key in [key for key in selected if '>' in key]:
                mainkey, subkey = key.split('>')
                if mainkey not in selected:
                    selected[mainkey] = {}
                selected[mainkey][subkey] = selected[key]
                del selected[key]  # 删除原始的嵌套键

            # 返回当前批次的数据
            yield selected